import dgl
import dgl.nn as dglnn
import torch
import pandas as pd
from utils import load_features
dataset_folder = '../data/'

def read_src_dst_weight(path,edge_type):
    src=[]
    dst=[]
    weight=[]
    data=pd.read_csv(path,header=None)
    #print("data:",data)
    for i in range(len(data[0])):
        src.append(data[0][i])
        dst.append(data[1][i])
        weight.append(data[2][i])
    #print("src:",src)
    #print("dst:",dst)
    #print("weight:",weight)
    return src, dst, weight


def get_graph():
    print('generating graph ...')
    #读取文件内容，获取edge相关信息
    src1,dst1,r2r_weight=read_src_dst_weight(dataset_folder+"edge/r2r/src1_dst1_r2r_weight.csv",'r-r')
    src2,dst2,r2s_weight=read_src_dst_weight(dataset_folder+"edge/r2s/src2_dst2_r2s_weight_0.csv",'r-s')
    src3,dst3,s2s_weight=read_src_dst_weight(dataset_folder+"edge/s2s/src3_dst3_s2s_weight_0.csv",'s-s')
    src4,dst4,s2i_weight=read_src_dst_weight(dataset_folder+"edge/s2i/src4_dst4_s2i_weight_0.csv",'s-i')
    src5,dst5,i2i_weight=read_src_dst_weight(dataset_folder+"edge/i2i/src5_dst5_i2i_weight_0.csv",'i-i')
    #src6,dst6,r2i_weight=read_src_dst_weight(dataset_folder+"edge/r2i/src6_dst6_r2i_weight_0.csv",'r-i')

    # nodes and edges
    graph = dgl.heterograph({#(起始点，结束点)
        ('recipe', 'r-r', 'recipe'): (torch.tensor(src1),torch.tensor(dst1)),
        ('recipe', 'r-s', 'step'): (torch.tensor(src2),torch.tensor(dst2)),
        ('step', 's-s', 'step'): (torch.tensor(src3),torch.tensor(dst3)),
        ('step', 's-i', 'ingredient'): (torch.tensor(src4),torch.tensor(dst4)),
        ('ingredient', 'i-s', 'step'): (torch.tensor(dst4),torch.tensor(src4)),
        ('ingredient', 'i-i', 'ingredient'): (torch.tensor(src5),torch.tensor(dst5)),
       # ('recipe', 'r-i', 'ingredient'): (torch.tensor(src6),torch.tensor(dst6))
    })



    # edge weight,,其实就是你通过余弦相似度或者tfidf计算出来的概率值，然后你去设定一个分解值（tfidf+cosine）例如0.5》0.5是，否则不是,i-i不需要edge weight,,然后s-s也不用吧，后续讨论
    graph.edges['r-r'].data['weight'] = torch.FloatTensor(r2r_weight)
    graph.edges['r-s'].data['weight'] = torch.FloatTensor(r2s_weight)
    graph.edges['s-s'].data['weight'] = torch.FloatTensor(s2s_weight)
    graph.edges['s-i'].data['weight'] = torch.FloatTensor(s2i_weight)
    graph.edges['i-s'].data['weight'] = torch.FloatTensor(s2i_weight)
    graph.edges['i-i'].data['weight'] = torch.FloatTensor(i2i_weight)

    # node features
    graph.nodes['recipe'].data['features'],graph.nodes['step'].data['features'],graph.nodes['ingredient'].data['features'] = load_features()
    return graph


